from sklearn.naive_bayes import MultinomialNB
from common import get_train_label
from scipy import sparse
import pickle
import os


# 训练贝叶斯模型
def trainBayesModel(name, path, header, sep=','):
    if not os.path.exists('../model/Bayes/%s_MNB_MODEL.pkl' % name):
        # 使用默认的配置对贝叶斯分类器进行初始化。
        mnb = MultinomialNB()
        train, label = get_train_label(path, header, sep=sep)
        train_fit = sparse.load_npz("../model/matrix/%s_TF_IDF_TRAIN_FIT_MATRIX.npz" % name)
        # 使用朴素贝叶斯分类器，对Tfidf后的训练样本进行参数学习。
        mnb.fit(train_fit, label)
        # 将训练好的贝叶斯模型保存到文件
        pickle.dump(mnb, open("../model/Bayes/%s_MNB_MODEL.pkl" % name, "wb"))
    else:
        print(name + "数据 Bayes 模型已进行训练，若需重新训练，删除文件即可")


if __name__ == '__main__':
    # 以微博数据训练贝叶斯模型
    trainBayesModel(name='weibo', path='../data/weibo_senti_100k.csv', header=['label', 'review'])
    print('微博数据Bayes模型训练完毕')
    # 以外卖数据训练贝叶斯模型
    trainBayesModel(name='waimai', path='../data/waimai_10k.csv', header=['label', 'review'])
    print('外卖数据Bayes模型训练完毕')
    # 以电商数据训练贝叶斯模型
    trainBayesModel(name='shopping', path='../data/online_shopping_10_cats.csv', header=['cat', 'label', 'review'])
    print('电商数据Bayes模型训练完毕')
    trainBayesModel(name='train', path='../data/train.txt', header=['label', 'review'], sep="\t")
    print('训练数据Bayes模型训练完毕')
